package binary_tree

func MaxDistance(root *Node) int {
	if root == nil {return 0}
	nodes := make([]*Node, 0)
	fn := func(n *Node) {
		nodes = append(nodes, n)
	}
	parentMap := make(map[*Node]*Node)
	parentMap[root] = nil
	FillParentMap(root, parentMap)
	GetAllNodes(root, fn)
	ans := 1
	for i:=0 ;i<len(nodes); i++ {
		for j:=i+1; j<len(nodes); j++ {
			ans = max(GetDistance(nodes[i], nodes[j], parentMap), ans)
		}
	}
	return ans
}

func GetDistance(n1 *Node, n2 *Node, parentMap map[*Node]*Node) int {
	parentSet := make(map[*Node]struct{})
	cur := n1
	for cur != nil {
		parentSet[cur] = struct{}{}
		cur = parentMap[cur]
	}
	cur = n2
	for cur != nil {
		if _, ok := parentSet[cur]; ok {
			break
		}
		cur = parentMap[cur]
	}
	ans := 1
	for n2 != cur {
		n2 = parentMap[n2]
		ans++
	}
	for n1 != cur {
		n1 = parentMap[n1]
		ans++
	}
	return ans
}

func GetAllNodes(root *Node, fn func(n *Node)) {
	if root == nil {return}
	fn(root)
	GetAllNodes(root.Left, fn)
	GetAllNodes(root.Right, fn)
}

func FillParentMap(root *Node, parentMap map[*Node]*Node){
	if root == nil {return}
	if root.Left != nil {
		parentMap[root.Left] = root
		FillParentMap(root.Left, parentMap)
	}
	if root.Right != nil {
		parentMap[root.Right] = root
		FillParentMap(root.Right, parentMap)
	}
}

type MaxDisInfo struct {
	MaxDis int
	height int
}

func MaxDistance1(root *Node) int {
	if root == nil {
		return 0
	}
	return MaxDisProcess(root).MaxDis
}

func MaxDisProcess(node *Node) MaxDisInfo {
	if node == nil {
		return MaxDisInfo{MaxDis:0, height: 0}
	}
	leftInfo := MaxDisProcess(node.Left)
	rightInfo := MaxDisProcess(node.Right)
	curHeight := max(leftInfo.height, rightInfo.height) + 1
	curMaxDistance := max(max(leftInfo.height+rightInfo.height+1, leftInfo.MaxDis), rightInfo.MaxDis)
	return MaxDisInfo{height:curHeight, MaxDis: curMaxDistance}
}